from fastapi import FastAPI, File, UploadFile, Form, Request
from fastapi.responses import JSONResponse, StreamingResponse, HTMLResponse
import io, os
from PIL import Image
import cv2
import numpy as np
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()

# 跨域配置（解决前端跨域问题）
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 挂载静态文件目录，访问同目录下的 index.html 等静态文件
app.mount("/static", StaticFiles(directory="."), name="static")

# 返回 index.html 的路由
@app.get("/")
async def read_index():
    with open("index.html", "r", encoding="utf-8") as f:
        content = f.read()
    return HTMLResponse(content=content)

# 获取 LoRA 列表
@app.get('/lora-list')
def get_lora_list():
    lora_dir = './loras'
    if not os.path.exists(lora_dir):
        return {'loras': []}
    loras = [f for f in os.listdir(lora_dir) if f.endswith('.safetensors')]
    return {'loras': loras}

# 留白检测接口
@app.post('/detect')
def detect(file: UploadFile = File(...)):
    img_bytes = file.file.read()
    img = Image.open(io.BytesIO(img_bytes)).convert('RGB')
    img_np = np.array(img)

    # 图像预处理：高斯模糊去噪，核大小调整为(7,7)增强去噪效果
    gray = cv2.GaussianBlur(img_np, (7, 7), 0)
    gray = cv2.cvtColor(gray, cv2.COLOR_RGB2GRAY)
    # 对比度增强，使用CLAHE（限制对比度自适应直方图均衡化），更柔和地增强对比度
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    gray = clahe.apply(gray)

    # 自适应阈值分割：调整块大小为9，C值为1，让阈值更灵敏地适配天空等区域
    mask = cv2.adaptiveThreshold(
        gray,
        255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY,
        9,  # 块大小减小，提升对局部区域的适配
        1    # C值减小，增加被判定为留白的区域范围
    )

    # 形态学操作：调整核大小为(5,5)，先膨胀填充更大空隙，再腐蚀还原
    kernel = np.ones((5, 5), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_DILATE, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_ERODE, kernel)

    # 过滤小面积轮廓：降低面积阈值到300，保留更多较小的留白细节
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for cnt in contours:
        if cv2.contourArea(cnt) < 300:
            cv2.drawContours(mask, [cnt], -1, 0, -1)

    # 生成并返回掩码图像
    mask_img = Image.fromarray(mask)
    buf = io.BytesIO()
    mask_img.save(buf, 'PNG')
    buf.seek(0)
    return StreamingResponse(buf, media_type='image/png')

# 生成接口（LoRA 可选）
@app.post('/generate')
def generate(
    file: UploadFile = File(...),
    mask_file: UploadFile = File(...),
    prompt: str = Form(...),
    negative_prompt: str = Form(''),
    lora: str = Form(None),
    guidance_scale: float = Form(7.5),
    num_inference_steps: int = Form(35),
    strength: float = Form(0.9)
):
    original = Image.open(file.file).convert('RGB')
    mask = Image.open(mask_file.file).convert('L')

    # 占位生成：直接返回原图+白色填充掩码示意
    result = original.copy()
    mask_np = np.array(mask)
    result_np = np.array(result)
    result_np[mask_np>128] = 255  # 将留白区域填白
    result_img = Image.fromarray(result_np)

    buf = io.BytesIO()
    result_img.save(buf, 'PNG')
    buf.seek(0)
    return StreamingResponse(buf, media_type='image/png')

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)